from PIL import Image, ImageDraw, ImageFont
import os
import platform

# -------------------------- 配置参数（可根据需求修改） --------------------------
# 字体路径：自动适配Windows/macOS/Linux系统，若自定义路径需替换下方对应值
SYSTEM = platform.system()
if SYSTEM == "Windows":
    # Windows默认字体路径（可替换为自己的字体路径，如"C:/Windows/Fonts/simhei.ttf"）
    FONT_PATH = "C:/Windows/Fonts/STXINGKA.TTF"
elif SYSTEM == "Darwin":  # macOS
    # macOS默认字体路径（可替换为"/Library/Fonts/Arial Unicode.ttf"等）
    FONT_PATH = "/Library/Fonts/Songti.ttc"
else:  # Linux（保留原路径，可根据发行版调整）
    FONT_PATH = "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc"

# 输入/输出目录（按你的路径保持不变）
INPUT_DIR = "../add_text/pomeranian_colors_annotated_final"
OUTPUT_DIR = "../add_text/pomeranian_colors_annotated_with_rare"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 需要添加“稀有”的图片文件名列表
RARE_FILENAMES = [
    "16_Blue.png", "17_Blue_Sable.png", "18_Blue_and_Tan.png", "19_Chocolate_Merle.png", "20_Beaver.png",
    "21_Black_and_Brindle.png", "22_Chocolate_Sable.png", "23_Lavender.png", "24_Blue_Brindle.png",
    "25_Merle_and_Tan.png"
]

# 文字样式配置（可自定义）
RARE_TEXT = "稀有"
RARE_COLOR = (255, 0, 0)  # 纯红色（可改为其他RGB值，如(220,0,0)深红外）
FONT_SIZE = 60  # 固定字体大小（可按需调整，如40、60）
DISTANCE_FROM_CORNER = 20  # 文字距离右上角的像素间距（越大越靠内）
ROTATION_ANGLE = 45  # 文字倾斜角度（固定45度，可改为30、60等）


# -------------------------- 核心工具函数 --------------------------
def load_font(size):
    """用你的字体路径逻辑加载字体，失败时用默认字体"""
    try:
        # 优先使用你配置的系统字体路径
        return ImageFont.truetype(FONT_PATH, size)
    except (IOError, OSError) as e:
        print(f"⚠️ 配置的字体路径加载失败（{str(e)}），将使用默认字体")
        # 字体加载失败时，降级使用PIL默认字体（确保不报错）
        return ImageFont.load_default()


def add_rotated_rare_text(input_path, output_path):
    """添加45度倾斜的“稀有”文字，确保完全在图片内"""
    try:
        # 打开已加蒙版的原图
        with Image.open(input_path).convert("RGB") as img:
            img_width, img_height = img.size
            # 加载字体（用你的配置逻辑）
            font = load_font(FONT_SIZE)

            # 1. 计算原始文字的宽高（用于创建临时画布）
            temp_draw = ImageDraw.Draw(Image.new('RGB', (1, 1)))
            text_bbox = temp_draw.textbbox((0, 0), RARE_TEXT, font=font)
            text_width = text_bbox[2] - text_bbox[0]
            text_height = text_bbox[3] - text_bbox[1]

            # 2. 创建临时透明画布（足够大，避免旋转裁剪文字）
            temp_canvas_size = max(text_width, text_height) + 20  # 预留边缘空间
            temp_canvas = Image.new('RGBA', (temp_canvas_size, temp_canvas_size), (0, 0, 0, 0))
            temp_canvas_draw = ImageDraw.Draw(temp_canvas)

            # 3. 在临时画布中心绘制“稀有”文字（确保旋转中心点正确）
            text_x = (temp_canvas_size - text_width) // 2
            text_y = (temp_canvas_size - text_height) // 2
            temp_canvas_draw.text((text_x, text_y), RARE_TEXT, font=font, fill=RARE_COLOR)

            # 4. 旋转临时画布（45度，expand=True确保旋转后尺寸足够）
            rotated_canvas = temp_canvas.rotate(ROTATION_ANGLE, expand=True)
            rotated_width, rotated_height = rotated_canvas.size

            # 5. 计算安全粘贴位置（确保文字不超出图片边界）
            # 基准位置：距离右上角DISTANCE_FROM_CORNER像素
            base_x = img_width - DISTANCE_FROM_CORNER - rotated_width // 2
            base_y = DISTANCE_FROM_CORNER - rotated_height // 2
            # 限制X坐标：不超出左边界（≥0）、不超出右边界（≤img_width - rotated_width）
            safe_x = max(0, min(base_x, img_width - rotated_width))
            # 限制Y坐标：不超出上边界（≥0）、不超出下边界（≤img_height - rotated_height）
            safe_y = max(0, min(base_y, img_height - rotated_height))

            # 6. 将旋转后的文字粘贴到原图（保留透明背景，不覆盖原图）
            img.paste(rotated_canvas, (int(safe_x), int(safe_y)), rotated_canvas)

            # 7. 保存最终结果
            img.save(output_path, quality=95)
            print(f"✅ 处理成功：{os.path.basename(output_path)}（位置：({safe_x:.0f},{safe_y:.0f})）")

    except Exception as e:
        print(f"❌ 处理失败 {os.path.basename(input_path)}：{str(e)}")


# -------------------------- 批量处理入口 --------------------------
if __name__ == "__main__":
    # 检查输入目录是否存在
    if not os.path.exists(INPUT_DIR):
        print(f"❌ 错误：输入目录「{INPUT_DIR}」不存在，请检查路径是否正确！")
        exit(1)

    # 遍历所有需要添加“稀有”文字的图片
    for filename in RARE_FILENAMES:
        input_img_path = os.path.join(INPUT_DIR, filename)
        output_img_path = os.path.join(OUTPUT_DIR, filename)

        # 检查图片文件是否存在
        if os.path.exists(input_img_path):
            add_rotated_rare_text(input_img_path, output_img_path)
        else:
            print(f"⚠️ 跳过：图片「{filename}」在输入目录中未找到")

    print(f"\n📌 所有图片处理完成！最终文件保存目录：{os.path.abspath(OUTPUT_DIR)}")